#! -*- coding: utf-8 -*-

import struct
import os
import time
import six
import codecs
import math
import logging
logging.basicConfig(level=logging.INFO, format=u'%(asctime)s - %(levelname)s - %(message)s')

"""
优化：
    1、使用了语言模型工具kenlm的count_ngrams程序来统计n-gram。
    由于kenlm是用C++写的，速度有保证，并且它还做了优化，所以对内存很友好。
    2、在第二次遍历词库以得到候选词的时候，使用了Trie树结构来加速搜索字符串是否出现过某个n-gram。
    *Trie树或者其变种，基本上是所有基于词典的分词工具的标配，就是因为它可以加快搜索字符串中是否出现过词典中的词。
    
凝固度：取对数之后就是我们所说的互信息熵
自由度：边界熵
"""


class Progress(object):
    """
    进度显示 (自己简单封装，比 tqdm 更可控一些)
    iterator: 可迭代的对象；
    period: 显示进度的周期；
    steps: iterator可迭代的总步数，相当于len(iterator)
    """
    def __init__(self, iterator, period=1, steps=None, desc=None):
        self.iterator = iterator
        self.period = period
        if hasattr(iterator, '__len__'):
            self.steps = len(iterator)
        else:
            self.steps = steps
        self.desc = desc
        if self.steps:
            self._format_ = u'%s/%s passed' %('%s', self.steps)
        else:
            self._format_ = u'%s passed'
        if self.desc:
            self._format_ = self.desc + ' - ' + self._format_
        self.logger = logging.getLogger()

    def __iter__(self):
        for i, j in enumerate(self.iterator):
            if (i + 1) % self.period == 0:
                self.logger.info(self._format_ % (i+1))
            yield j


class KenlmNgrams(object):
    """
    加载 Kenlm 的 n-gram统计结果 (读取ngrams.txt文件中的统计结果)
        vocab_file: Kenlm统计出来的词(字)表；
        ngram_file: Kenlm统计出来的 n-gram表；
        order: 统计 n-gram 时设置的 n，必须跟 ngram_file 对应；
        min_count: 自行设置的截断频数（词频）
    """
    def __init__(self, vocab_file, ngram_file, order, min_count):
        self.vocab_file = vocab_file
        self.ngram_file = ngram_file
        self.order = order
        self.min_count = min_count
        self.read_chars()
        self.read_ngrams()

    def read_chars(self):
        f = open(self.vocab_file)
        chars = f.read()
        f.close()
        chars = chars.split('\x00')
        self.chars = [i.decode('utf-8') if six.PY2 else i for i in chars]

    def read_ngrams(self):
        """
        读取思路参考 https://github.com/kpu/kenlm/issues/201
        """
        self.ngrams = [{} for _ in range(self.order)]
        self.total = 0
        # Vocabulary ids as 4-byte integers. Since you've selected order 3, there are 3 of them.
        # With an 8-byte integer-counter, so in this case each record is 3*4 + 8 = 20 bytes
        size_per_item = self.order * 4 + 8  # ngrams.txt中每条记录的字节数
        f = open(self.ngram_file, 'rb')
        filedata = f.read()
        filesize = f.tell()  # 获取文件中的记录总数
        f.close()
        # 需要注意的是，这里没有按行提取，是因为 ngrams.txt中二进制不是按行存储的
        for i in Progress(iterator=range(0, filesize, size_per_item), period=100000, desc=u'loading ngrams'):
            s = filedata[i: i+size_per_item]  # 获取每条记录
            # 取出[2-gram, 3-gram, 4-gram]中的2-gram的词频
            n = self.unpack('l', s[-8:])  # 去掉 integer-counter的 8个字节后解包，“l” 表示按照 “无字符长整型” 格式
            if n >= self.min_count:
                self.total += n
                # "i" 表示int, 所以这里取出的是 3个order不同 gram 的词频: [2-gram, 3-gram, 4-gram]
                c = [self.unpack('i', s[j*4: (j+1)*4]) for j in range(self.order)]
                c = ''.join([self.chars[num] for num in c if num > 2])
                for j in range(len(c)):
                    self.ngrams[j][c[:j+1]] = self.ngrams[j].get(c[:j+1], 0) + n

    def unpack(self, t, s):
        # struct.unpack() 方法的使用详见: https://blog.csdn.net/weiwangchao_/article/details/80395941
        return struct.unpack(t, s)[0]  # 取出[2-gram, 3-gram, 4-gram]中的2-gram的词频


def write_corpus(texts, filename):
    """
    将语料写到文件中，词与词(字与字) 之间用空格隔开
    """
    with codecs.open(filename, 'w', encoding='utf-8') as f:
        for s in Progress(iterator=texts, period=10000, desc=u'exporting corpus'):
            s = ' '.join(s) + '\n'
            f.write(s)


def count_ngrams(corpus_file, order, vocab_file, ngram_file, memory=0.5):
    """
    通过os.system调用Kenlm的count_ngrams来统计频数
        - memory: 占用内存比例，理论上不能超过可用内存比例
        - order: n-gram的层级, 表示需要2-gram, 3-gram, 4-gram
    """
    done = os.system(
        './count_ngrams -o %s --memory=%d%% --write_vocab_list %s <%s >%s' % (order, memory * 100, vocab_file, corpus_file, ngram_file)
    )
    if done != 0:
        raise ValueError('Failed to count n-grams by KenLM.')


def filter_ngrams(ngrams, total, min_pmi=1):
    """
    通过 “互信息” 过滤 ngrams，只保留 “结实” (n-gram凝固度高) 的 ngram。
    """
    order = len(ngrams)
    if hasattr(min_pmi, '__iter__'):
        min_pmi = list(min_pmi)  # 如果ngram中的n是多个值，对应的 “凝固度” min_pmi 也是多个
    else:
        min_pmi = [min_pmi] * order  # 如果ngram中的n是单个值，
    output_ngrams = set()
    total = float(total)
    for i in range(order-1, 0, -1):
        for w, v in ngrams[i].items():

            pmi = min([total * v / (ngrams[j].get(w[:j+1], total) * ngrams[i-j-1].get(w[j+1:], total)) for j in range(i)])
            if math.log(pmi) >= min_pmi[i]:
                output_ngrams.add(w)

    return output_ngrams


class SimpleTrie(object):
    """
    通过 Trie树结构，来搜索 n-grams 组成的连续片段
    """
    def __init__(self):
        self.dic = {}
        self.end = True

    def add_word(self, word):
        _ = self.dic
        for c in word:
            if c not in _:
                _[c] = {}
            _ = _[c]
        _[self.end] = word

    def tokenize(self, sent):
        # 通过 "最长联接" 的方式来对句子进行分词
        result = []
        start, end = 0, 1
        for i, c1 in enumerate(sent):
            _ = self.dic
            if i == end:
                result.append(sent[start: end])
                start, end = i, i+1
            for j, c2 in enumerate(sent[i:]):
                if c2 in _:
                    _ = _[c2]
                    if self.end in _:
                        if i + j + 1 > end:
                            end = i + j + 1
                else:
                    break
        result.append(sent[start: end])
        return result


def filter_vocab(candidates, ngrams, order):
    """
    通过与 n-grams 对比，排除可能的 “不牢固” 词汇 (回溯)
    """
    result = {}
    for i, j in candidates.items():
        if len(i) < 3:
            result[i] = j
        elif len(i) <= order and i in ngrams:
            result[i] = j
        elif len(i) > order:
            flag = True
            for k in range(len(i) + 1 - order):
                if i[k: k+order] not in ngrams:
                    flag = False
            if flag:
                result[i] = j
    return result

    # def text_generator():
    #     """
    #     预料生成器，并初步预处理。通过yield一句一句的给出来
    #     """
    #     d = open('../data/ftp_Title_Content_Corpus20200726.csv', encoding='utf-8').read()
    #     d = d.replace(u'\u30000', ' ').strip()
    #     yield re.sub(u'[^\u4e00-\u9fa50-9a-zA-Z ]+', '\n', d)

    # def generate_corpus(ori_corpus_file, corpus_file):
    #     ori_f = open(ori_corpus_file, 'r', encoding='utf-8')
    #     corpus_f = open(corpus_file, 'w', encoding='utf-8')
    #     line = ori_f.readline()
    #     while line:
    #         content = line.replace(u'\u30000', ' ').strip()
    #         texts = re.sub(u'[^\u4e00-\u9fa50-9a-zA-Z ]+', '\n', content)
    #         for s in Progress(texts, 10000, desc=u'exporting corpus'):
    #             s = ' '.join(s) + '\n'
    #             corpus_f.write(s)
    #
    #         line = ori_f.readline()
    #
    #     ori_f.close()
    #     corpus_f.close()


# ======= 算法构建完毕，下面开始执行完整的构建词库流程 =======
if __name__ == '__main__':

    import re
    import glob

    # 语料生成器，并且初步预处理语料(只保留 “中文、数字、大小写字母 ”)
    # 生成器具体含义: 逐句地把文本yield出来
    def text_generator():
        txts = glob.glob('../data/origin/*.txt')
        for txt in txts:
            d = codecs.open(txt, encoding='utf-8').read()
            d = d.replace(u'\u3000', ' ').strip()  # \u3000是unicode中的全角空格，这里是将它变为代码常用的半角空格（\u0020）
            yield re.sub(u'[^\u4e00-\u9fa50-9a-zA-Z ]+', '\n', d)  # 只保留 “中文、数字、大小写字母、半角空格”

    # min_count = 32  # 频数过滤的阈值
    min_count = 2  # 频数过滤的阈值
    order = 4  # n-gram的层级 n
    memory = 0.5  # memory是占用内存比例，理论上不能超过可用内存比例

    # ori_corpus_file = '../data/test.txt'  # 测试样例文件名

    # corpus_file = '../data/corpus.txt'  # 语料保存的文件名
    corpus_file = '../data/corpus0.txt'  # 语料保存的文件名
    # vocab_file = '../data/chars.txt'  # 字符集保存的文件名
    vocab_file = '../data/chars0.txt'  # 字符集保存的文件名
    # ngram_file = '../data/ngrams.txt'  # ngram集保存的文件名
    ngram_file = '../data/ngrams0.txt'  # ngram集保存的文件名
    # output_file = '../data/vocab.txt'  # 导出的 "新词" 词表文件名
    output_file = '../data/red_vocab.txt'  # 导出的 "新词" 词表文件名

    # kenlm需要一个以空格分词的、纯文本格式的语料作为输入
    t1 = time.time()
    # 将语料写到文件中，词与词(字与字) 之间用空格隔开
    write_corpus(text_generator(), corpus_file)
    t2 = time.time()
    print("生成语料corpus共耗时：{}".format(t2-t1))

    # count_ngrams() 就是调用 kenlm（基于c++, 速度快）的count_ngrams程序来统计n-gram，生成二进制文件
    t1 = time.time()
    count_ngrams(corpus_file, order, vocab_file, ngram_file, memory)  # 用Kenlm统计ngram
    t2 = time.time()
    print("生成n-grams结果共耗时：{}".format(t2 - t1))

    # KenlmNgrams 读取二进制文件
    t1 = time.time()
    ngrams = KenlmNgrams(vocab_file, ngram_file, order, min_count)  # 加载ngrams
    t2 = time.time()
    print("生成二进制文件共耗时：{}".format(t2 - t1))

    # [0, 2, 4, 6]是互信息的阈值，其中第一个0无意义，仅填充用，而 2, 4, 6分别是 2gram、3gram、4gram 的互信息阈值
    t1 = time.time()
    # print("ngrams.ngrams: ", ngrams.ngrams)
    print("ngrams.total: ", ngrams.total)
    # ngrams = filter_ngrams(ngrams.ngrams, ngrams.total, [0, 4, 6, 8])  # 过滤 ngram, 增大凝固度
    # ngrams = filter_ngrams(ngrams.ngrams, ngrams.total, [0, 3, 5, 7])  # 过滤 ngram, 增大凝固度
    ngrams = filter_ngrams(ngrams.ngrams, ngrams.total, [0, 2, 4, 6])  # 过滤 ngram
    # ngrams = filter_ngrams(ngrams.ngrams, ngrams.total, [0, 1, 3, 5])  # 过滤 ngram, 降低凝固度
    t2 = time.time()
    print("准备工作已经完成，共耗时：{}".format(t2 - t1))

    # =================准备工作已完成=========================

    # 构建一个n-gram的Trie树，然后用这个Trie树就可以做一个基本的 “预分词” ，有点像最大匹配，由 n-gram片段连接成尽可能长的候选词
    ngtrie = SimpleTrie()  # 构建一个ngram的Trie树
    t1 = time.time()
    for w in Progress(iterator=ngrams, period=100000, desc=u'build ngram trie'):
        _ = ngtrie.add_word(w)
    t2 = time.time()
    print("Trie树构建完毕，共耗时：{}".format(t2 - t1))

    # 构建候选词库
    candidates = {}
    t1 = time.time()
    for t in Progress(iterator=text_generator(), period=1000, desc='discovering words'):
        for w in ngtrie.tokenize(t):  # 预分词： 通过 "最长联接" 的方式来对句子进行分词
            candidates[w] = candidates.get(w, 0) + 1
    t2 = time.time()
    print("候选词构建完毕，共耗时：{}".format(t2 - t1))

    # 频数过滤：过滤掉词频较小的候选词
    t1 = time.time()
    candidates = {i: j for i, j in candidates.items() if j >= min_count}
    t2 = time.time()
    print("候选词频数过滤完毕，共耗时：{}".format(t2 - t1))

    # 互信息过滤(回溯)
    t1 = time.time()
    candidates = filter_vocab(candidates, ngrams, order)
    t2 = time.time()
    print("候选词互信息过滤完毕，共耗时：{}".format(t2 - t1))
    # 输出结果文件
    with codecs.open(output_file, 'w', encoding='utf-8') as f:
        for i, j in sorted(candidates.items(), key=lambda s: -s[1]):
            if len(i) > 2:  # 只保留长度大于1的
                s = '%s %s\n' % (i, j)
                f.write(s)

    # str = "拓维信息(002261)总经理减持200万股,拓维信息(002261)2014年1月1日晚间公告，公司副董事长、总经理宋鹰因个人资金需求，" \
    #       "于2013年12月30日通过深圳证券交易所大宗交易方式减持其持有的公司无限售条件流通股共计2,000,000股，" \
    #       "占公司总股本的0.71%。。本次减持后，宋鹰持有公司股份5829.99万股，占公司总股本的20.56%。,2014-01-01T00:00:00+0800," \
    #       "002261_SZ,拓维信息,2014-01-01T00:00:00+0800"
    #
    # # 对除（大小写字母、空格、汉字 ）之外的符号做切分处理
    # res = re.sub(u'[^\u4e00-\u9fa50-9a-zA-Z ]+', '\n', str)
    #
    # print(res)